import datetime

from matplotlib import pyplot as plt
from torch import nn, Tensor
import torch


def read_data(path):
    date_list = []
    date_inter_list = []
    disk_usage = []
    with open(path, 'r', encoding='UTF-8') as input_data:
        # 逐条读取记录
        i = 0
        for line in input_data:
            if i == 0:
                i = 1
                continue
            line_list = line.split()

            data_date = datetime.datetime.strptime(line_list[0], "%Y-%m-%d").date()
            date_list.append(data_date)
            inter_date = (data_date - start_date).days
            date_inter_list.append(inter_date)
            disk_usage.append(float(line_list[1]))
    return date_inter_list, disk_usage, date_list


start_date_str = '2024-05-13'
start_date = datetime.datetime.strptime(start_date_str, "%Y-%m-%d").date()

end_date_str = '2024-07-01'
end_date = datetime.datetime.strptime(end_date_str, "%Y-%m-%d").date()

path = "./disk_usage"
date_inter_list, disk_usage, date_list = read_data(path)

predict_date_list = []
predict_date_inter_list = []
actul_last_date = date_list[-1]
for i in range((end_date - actul_last_date).days):
    tmp_date = actul_last_date + datetime.timedelta(days=i)
    predict_date_list.append(tmp_date)
    predict_date_inter_list.append((tmp_date - start_date).days)

x = Tensor(date_inter_list)
y = Tensor(disk_usage)

x = x.reshape(-1, 1)
y = y.reshape(-1, 1)

lr = 0.001
epochs = 300


class ClickHouseDiskUsage(nn.Module):
    def __init__(self):
        super(ClickHouseDiskUsage, self).__init__()  # 继承父类init的参数
        self.linear = nn.Linear(1, 1)

    def forward(self, x):
        out = self.linear(x)
        return out


model = ClickHouseDiskUsage()

loss_fn = nn.MSELoss()

optimizer = torch.optim.SGD(model.parameters(), lr=lr)

for epoch in range(epochs):
    y_predict = model(x)
    loss = loss_fn(y, y_predict)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if (epoch + 1) % 1000 == 0:
        print("epoch-{}: loss = {}".format(epoch + 1, loss))

model.eval()  #设置模型为评估模式，即预测模式
predict = model(x)
predict = predict.data.numpy()
plt.figure(num=1, figsize=(10, 5))
plt.xlabel('date')
plt.ylabel('GiB')
plt.scatter(date_list, y.data.numpy(), c="r")
plt.plot(date_list, predict)
plt.grid()
plt.show()

plt.figure(num=2, figsize=(10, 5))
plt.xlabel('date')
plt.ylabel('GiB')
plt.xlim((date_list[0], end_date))
# plt.ylim((0, 200))
plt.scatter(date_list, y.data.numpy(), c="r")
plt.plot(date_list, predict)
predict2 = model(Tensor(predict_date_inter_list).reshape(-1, 1))
predict2 = predict2.data.numpy()
plt.plot(predict_date_list, predict2, c="g", marker='o', markeredgecolor='b', markersize='5')
plt.grid()
for i, j in zip(predict_date_list, predict2):
    plt.text(i, j-2, '(%s,%.1f)' % (i,j[0]))
plt.show()
